#!/usr/bin/env python

import psycopg2
import psycopg2.pool
import matplotlib
import concurrent.futures
import traceback
matplotlib.use('Agg')
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt 
plt.style.use('ggplot')

dbcon_pool = psycopg2.pool.ThreadedConnectionPool(8, 40, 
        "dbname=trending_topic_alpha user=atropos host='/var/run/postgresql/'") 


def gntsp():
    """
    Query for time span
    """
    db = dbcon_pool.getconn()
    cur = db.cursor()
    cur.execute("select min(timestamp) from topic_reads where timestamp like '%%01';")
    start = cur.fetchone()[0]
    cur.execute("select max(timestamp) from topic_reads where timestamp like '%%00';")
    end = cur.fetchone()[0]
    cur.close()
    dbcon_pool.putconn(db)
    return (start, end)

def gndtp(date, db):
    """
    Return a list of topics of the specified date.
    The date should be of the form like '20160908'
    """
    cur = db.cursor()
    cur.execute("""select distinct id from topic_reads where timestamp like
            '%s%%'; """, (date,)) 
    tids = [tid[0] for tid in cur.fetchall()]
    cur.close()
    return tids 

def gnddlt(date, tids, db):
    """
    Return the top K topic reads
    """
    #print(date, end=" ")
    cur = db.cursor()
    reads = []
    for tid in tids:
        #print(".", end="", flush=True)
        cur.execute("""select reads from topic_reads where id = %s and 
                timestamp like '%s%%'; """, (tid, date))
        rdsum = 0
        for row in cur.fetchall():
            rdsum += row[0]
        cur.execute("""select reads from topic_reads where id = %s and
                timestamp like '%s00';""", (tid, (date + 1)))
        rdtmp = cur.fetchone()
        if rdtmp:
            rdsum += rdtmp[0]
        cur.execute("""select reads from topic_reads where id = %s and
                timestamp like '%s00';""", (tid, date))
        rdtmp = cur.fetchone()
        if rdtmp:
            rdsum -= rdtmp[0]
        reads.append(rdsum)
    #reads = sorted(reads)
    cur.close()
    return reads

def gntkgh(tprds, start, end, topk, sumk):
    """
    Plot some basic analysis graphs
    """
    if start == end:
        print("Time window is too tight to plot")
    else:
        tops = ['Top' + str(i) for i in range(topk, 0, -1)]
        total_per_day = [sum(day) for day in tprds]

        tprds = [day[-sumk:] for day in tprds]
        total_for_sumk = [sum(day) for day in tprds]
        avg_ptrds = int(sum(total_for_sumk) / len(total_for_sumk))
        tprds = [day[-topk:] for day in tprds]

        total_number = 1000

        first_k_tprds = [day[-total_number:] for day in tprds]
        first_k_tprds = [sum(k) for k in first_k_tprds]
        first_k_tprds = np.array(first_k_tprds, dtype=np.uint)

        tprds = np.array(tprds, dtype=np.float)
        for i in range(0,len(tprds)):
            tprds[i] = [reads_per_topic/total_for_sumk[i] for reads_per_topic in tprds[i]]
        print("Time window: %s ~ %s(Not included)" % (start, end))
        print("\tTotal number of topic reads per day:")
        print("\t", total_per_day)
        print("\tTotal number of top ", sumk, " topic reads per day:")
        print("\t", total_for_sumk)
        print("\tAverage number of topic reads per day: ", avg_ptrds)
        print("\tTotal number of top ", total_number, " topic reads per day:")
        print("\t", first_k_tprds)
        print("\n\tTop %s topic reads" % (topk,))
        i = 0
        for rds in tprds:
            print("\t", end='', flush=True)
            for rd in rds:
                print("%.3f" % (rd,), end=" ", flush=True)
            print("Sum: %.2f%%" % (sum(rds)*100), "Top", sumk, "/Total: %.2f%%" %
                    ((total_for_sumk[i]/total_per_day[i])*100))
            i += 1
        df = pd.DataFrame(tprds, index = pd.date_range(start, str(int(end)-1)), columns = tops)
        plt.figure()
        ax = df.plot()
        ax.get_figure().savefig('percentage.png')

        df = pd.DataFrame(first_k_tprds, index = pd.date_range(start, str(int(end)-1)), columns = ['Total Reads'])
        plt.figure()
        ax = df.plot()
        ax.get_figure().savefig('total.png')


def plgnddlt(date):
    # Caculate top K topic reads for a day
    global dbcon_pool
    db = dbcon_pool.getconn()
    tids = gndtp(date, db)
    reads = []
    with concurrent.futures.ThreadPoolExecutor() as executor:
        futures = dict()
        i = 0
        step = 100 
        while i < len(tids):
            if(i+step >= len(tids)):
                futures[executor.submit(gnddlt,date, tids[i:], db)] = i
            else:
                futures[executor.submit(gnddlt,date, tids[i:i+step], db)] = i
            i += step
        for future in concurrent.futures.as_completed(futures):
            try:
                reads.extend(future.result())
            except Exception:
                traceback.print_exc()
    reads = sorted(reads)
    dbcon_pool.putconn(db)
    return reads

def genDateSpan(dt_from, dt_to):
    """
    Generate a list of date stamps from dt_from to 
    dt_to excluded.
    """
    def _numDays(d):
        if d in [1, 3, 5, 7, 8, 10, 12]:
            return 31
        elif d is 2:
            return 28
        else:
            return 30

    days = []
    mn_from = (dt_from // 100) % 100
    mn_to = (dt_to // 100) % 100
    mn_delta = mn_to - mn_from
    day_from = dt_from % 100
    day_to = dt_to % 100

    for m in range(mn_from, mn_to+1):
        if m is mn_from:
            if mn_from is mn_to:
                days.extend([i for i in range(dt_from, dt_to+1)])
            else:
                days.extend([dt_from+d for d in range(_numDays(m)-day_from+1)])
        elif m is mn_to:
            days.extend([2016*10000 + m*100 + d for d in range(1, mn_to+1)])
        else:
            days.extend([2016*10000 + m*100 + d for d in range(1, _numDays(m)+1)])
    days.pop(-1)
    return days




if __name__ == '__main__':
    topk = 10
    sumk = 50 #Top k users we used to caculate the number of 
              #reads for entire platform
    #start, end = gntsp()
    #if not start or not end:
    #    import sys
    #    sys.exit()
    #start = start[:-2]
    #end = end[:-2]
    start = '20161015'
    end = '20161016'
    print("From ", start, " to ", end)
    tprds = dict()
    with concurrent.futures.ThreadPoolExecutor() as executor:
        futures = dict()
        for date in genDateSpan(int(start), int(end)):
            futures[executor.submit(plgnddlt, date)] = date
        for future in concurrent.futures.as_completed(futures):
            try:
                tprds[futures[future]] = future.result()
            except Exception:
                traceback.print_exc()
    dbcon_pool.closeall()
    tprds_array = []
    for date in sorted(tprds.keys()):
        tprds_array.append(tprds[date])
    gntkgh(tprds_array, start, end, topk, sumk)
    

